import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import random
from Utils.tensor_ops import _2tensor
from collections import namedtuple
from Utils.tensor_ops import pt_inf
from torch.distributions.categorical import Categorical

# 定义经验回放缓冲区的数据结构
# Transition: 用于存储(s, a, s', r, done)元组，这是强化学习中的基本经验单元
Transition = namedtuple('Transition', ('state', 'action', 'next_state', 'reward', 'done'))


class ReplayBuffer:
    """
    经验回放缓冲区 - 强化学习的核心组件之一

    关键特性：
    1. 循环缓冲区设计：当缓冲区满时，覆盖最旧的数据
    2. 随机采样：打破时间相关性，提高样本独立性
    3. 固定容量：防止内存无限增长，促进旧数据遗忘

    设计选择：
    - 使用namedtuple确保数据结构清晰
    - 模运算实现高效的循环覆盖
    - 预分配None避免频繁内存分配
    """
    def __init__(self, capacity):
        self.capacity = capacity  # 缓冲区最大容量
        self.memory = []         # 实际存储数据的列表
        self.position = 0        # 当前写入位置指针

    def push(self, *args):
        """
        存储一个transition到缓冲区

        实现细节：
        1. 如果缓冲区未满，添加新元素
        2. 如果缓冲区已满，覆盖最旧元素
        3. 使用模运算实现循环写入
        """
        if len(self.memory) < self.capacity:
            # 缓冲区未满，添加新元素
            self.memory.append(None)
        # 循环覆盖：position在[0, capacity-1]范围内循环
        self.memory[self.position] = Transition(*args)
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        """
        随机采样一批transitions

        重要意义：
        - 打破时间序列相关性，满足i.i.d.假设
        - 提高样本利用率，每条样本可被多次使用
        - 减少学习方差，提高训练稳定性
        """
        return random.sample(self.memory, batch_size)

    def __len__(self):
        """返回当前缓冲区中的样本数量"""
        return len(self.memory)



# 定义Q网络
class QNetwork(nn.Module):
    """
    单个智能体的Q网络 - 基于深度Q网络(DQN)的架构

    网络架构分析：
    - 输入层到隐藏层1：input_dim -> 512 (特征提取和初步抽象)
    - 隐藏层1到隐藏层2：512 -> 256 (高层次特征学习)
    - 隐藏层2到输出层：256 -> output_dim (Q值估计)

    设计选择：
    1. 三层全连接网络：平衡表达能力和计算效率
    2. ReLU激活函数：解决梯度消失，加速收敛
    3. 线性输出层：直接输出Q值，不需要激活函数

    技术细节：
    - 每个智能体都有独立的Q网络
    - 网络参数独立学习，实现去中心化决策
    - 输出维度等于动作空间大小
    """
    def __init__(self, input_dim, output_dim):
        super(QNetwork, self).__init__()
        # 特征提取层：将状态映射到高维特征空间
        self.fc1 = nn.Linear(input_dim, 512)
        # 抽象表示层：学习状态的高级特征
        self.fc2 = nn.Linear(512, 256)
        # Q值输出层：输出每个动作的Q值
        self.fc3 = nn.Linear(256, output_dim)

    def forward(self, x):
        """
        前向传播过程

        输入：x - 状态张量，shape=[batch_size, input_dim]
        输出：Q值张量，shape=[batch_size, output_dim]

        数据流：
        1. 第一层：线性变换 + ReLU激活
        2. 第二层：线性变换 + ReLU激活
        3. 第三层：线性变换输出Q值
        """
        x = torch.relu(self.fc1(x))  # 特征提取和初步抽象
        x = torch.relu(self.fc2(x))  # 高级特征学习
        x = self.fc3(x)              # Q值估计（无激活函数）
        return x


# QMIX算法实现 - 多智能体强化学习的核心算法
class QMIX(nn.Module):
    """
    QMIX (Q-value Mixing Network) 算法实现

    算法核心思想：
    1. 去中心化执行：每个智能体独立决策
    2. 中心化训练：利用全局信息进行训练
    3. Q值分解：将全局Q值分解为个体Q值
    4. 单调性约束：确保全局Q值与个体Q值的一致性

    架构特点：
    - 每个智能体拥有独立的Q网络
    - 共享经验回放缓冲区
    - 独立的优化器进行参数更新

    参数说明：
    - num_agents: 智能体数量
    - state_dim: 每个智能体的状态维度
    - action_space: 动作空间大小
    - device: 计算设备(CPU/GPU)
    """
    def __init__(self, num_agents, state_dim, action_space, device):
        super().__init__()
        self.num_agents = num_agents          # 智能体数量
        self.state_dim = state_dim            # 每个智能体的状态维度
        self.action_space = action_space      # 动作空间大小
        self.device = device                  # 计算设备
        self.num_thread = 0                   # 线程数（并行环境数量）

        # 为每个智能体创建独立的Q网络
        # ModuleList确保PyTorch能够正确跟踪所有网络的参数
        self.q_networks = torch.nn.ModuleList(
            modules=[QNetwork(self.state_dim, self.action_space) for _ in range(self.num_agents)])

        # 共享的经验回放缓冲区
        # 容量10000：平衡记忆能力和计算效率
        self.replay_buffer = ReplayBuffer(capacity=10000)

        # 为每个Q网络创建独立的Adam优化器
        # 学习率0.001：常用的默认值，平衡收敛速度和稳定性
        self.optimizers = [optim.Adam(q_net.parameters(), lr=0.001) for q_net in self.q_networks]
        # 优化器字典：方便按索引访问各个优化器
        self.optimizers_dict = {'optimizer' + str(index + 1): optimizer for index, optimizer in
                                enumerate(self.optimizers)}

    def select_action(self, states, test_mode, avail_act, epsilon=0.1):
        """
        动作选择策略 - 支持多线程并行环境

        算法流程：
        1. 遍历每个线程（环境实例）
        2. 对每个智能体计算Q值
        3. 处理可用动作约束
        4. 根据模式选择动作（探索vs利用）

        关键技术点：
        - epsilon-greedy探索策略
        - 可用动作掩码处理
        - 分类分布采样
        - 张量拼接优化

        参数：
        - states: 状态张量，shape=[num_threads, num_agents, state_dim]
        - test_mode: 是否为测试模式（无探索）
        - avail_act: 可用动作掩码
        - epsilon: 探索概率

        返回：
        - actions: 动作数组，shape=[num_threads, num_agents]
        """
        flag = True
        q_value = None
        self.num_thread = states.shape[0]  # 线程数
        action_list = torch.empty(1, self.num_agents).to(self.device)

        # 遍历每个线程（并行环境）
        for i in range(self.num_thread):
            # 重塑状态张量：分离各智能体的状态
            state_thread = states[i].reshape(self.num_agents, self.state_dim)
            # 处理可用动作掩码
            avail_act_thread = torch.tensor(avail_act[i],dtype=torch.int64).to(self.device)

            # 遍历每个智能体，计算Q值
            for q_network, state in zip(self.q_networks, state_thread):
                # 状态预处理：转换为张量并处理NaN值
                state_tensor = torch.tensor(state, dtype=torch.float32).to(self.device)
                state_tensor = torch.nan_to_num_(state_tensor, 0).to(self.device)

                # epsilon-greedy探索策略
                if random.random() < 0.01:  # 1%概率随机探索
                    action = random.randint(0, self.action_space - 1)
                    q_values = torch.zeros(self.action_space).to(self.device)
                    q_values[action] = 1
                else:  # 99%概率利用当前策略
                    with torch.no_grad():
                        q_values = q_network(state_tensor).to(self.device)

                # 拼接所有智能体的Q值
                if flag:
                    q_value = q_values.unsqueeze(0)
                    flag = False
                else:
                    q_value = torch.cat((q_value, q_values.unsqueeze(0)))

            # 应用可用动作掩码：将不可用动作的Q值设为负无穷
            if avail_act is not None:
                q_value = torch.where(avail_act_thread > 0, q_value, -pt_inf()).to(self.device)

            # 创建分类分布用于动作采样
            act_dist = Categorical(logits=q_value.unsqueeze(0))
            if not test_mode:
                # 训练模式：从概率分布中采样（探索）
                actions = act_dist.sample()
            else:
                # 测试模式：选择概率最高的动作（利用）
                actions = torch.argmax(act_dist.probs, dim=2)

            # 累积所有线程的动作
            if i == 0:
                action_list = actions
            else:
                action_list = torch.cat((action_list, actions), 0)

            # 重置临时变量
            q_value = None
            flag = True

        return action_list.cpu().numpy()

    def evaluate_actions(self, *args, **kargs):
        """评估动作的包装函数，直接调用更新函数"""
        return self._update(*args, **kargs)

    def _update(self, state_batch, action_batch, avail_act_batch, next_state_batch, reward_batch, done_batch, gamma,
                max_grad_norm):
        """
        QMIX的核心学习更新函数 - 实现独立Q-learning算法

        算法流程：
        1. 对每个智能体独立训练
        2. 计算目标Q值（Bellman方程）
        3. 计算时序差分(TD)误差
        4. 反向传播更新网络参数

        关键算法细节：
        - 目标网络：使用当前网络估计目标值（Double DQN的简化版）
        - 梯度裁剪：防止梯度爆炸
        - L1损失：对异常值比L2损失更鲁棒
        - 独立学习：每个智能体独立更新自己的Q网络

        数学基础：
        Q_target = r + γ * max_a' Q(s', a') * (1 - done)
        Loss = |Q(s,a) - Q_target|_1

        参数：
        - state_batch: 状态批次，shape=[batch_size, num_agents, time_steps, state_dim]
        - action_batch: 动作批次，shape=[batch_size, num_agents, time_steps]
        - avail_act_batch: 可用动作批次
        - next_state_batch: 下一状态批次
        - reward_batch: 奖励批次
        - done_batch: 终止标志批次
        - gamma: 折扣因子
        - max_grad_norm: 最大梯度范数
        """
        # 注释掉的代码：经验回放采样逻辑
        # if len(self.replay_buffer) < batch_size:
        #     return
        # transitions = self.replay_buffer.sample(batch_size)
        # batch = Transition(*zip(*transitions))

        loss_list = []

        # 遍历每个智能体进行独立学习
        for i in range(self.num_agents):
            # 数据预处理：提取当前智能体的数据
            state = state_batch[:, i, :, :].reshape(state_batch.shape[0], -1)
            next_state = next_state_batch[:, i, :, :].reshape(next_state_batch.shape[0], -1)
            # 处理NaN值：将其替换为0
            state = torch.nan_to_num_(state, 0)
            next_state = torch.nan_to_num_(next_state, 0)
            action = action_batch[:, i].to(torch.int64)
            avail_act = avail_act_batch[:, i, :]
            reward = reward_batch[:, i]
            done = done_batch[:, i]

            # 目标Q值计算（使用no_grad避免梯度传播）
            with torch.no_grad():
                # 计算下一状态的最大Q值
                q_values_next = self.q_networks[i](next_state)
                # 应用可用动作掩码
                if avail_act is not None:
                    q_values_next = torch.where(avail_act > 0, q_values_next, -pt_inf())
                # 取最大值（贪婪策略）
                q_values_next = q_values_next.max(1)[0]
                # Bellman方程：Q_target = r + γ * max_a' Q(s', a') * (1 - done)
                q_targets = reward + gamma * q_values_next * ~done

            # 当前状态Q值计算
            q_values = self.q_networks[i](state)
            # 计算TD误差：使用L1损失（Huber损失的特例）
            # gather(1, action.unsqueeze(1))：选择执行动作的Q值
            loss = nn.functional.l1_loss(q_values.gather(1, action.unsqueeze(1)), q_targets.unsqueeze(1))
            loss_list.append(loss.item())

            # 反向传播和参数更新
            self.optimizers[i].zero_grad()  # 清空梯度
            loss.backward()  # 反向传播计算梯度
            # 梯度裁剪：防止梯度爆炸，提高训练稳定性
            nn.utils.clip_grad_norm_(self.parameters(), max_grad_norm)
            self.optimizers[i].step()  # 更新网络参数

        # 计算平均损失（用于监控训练过程）
        mean_loss = np.array(loss_list).sum() / (state_batch.shape[0] * state_batch.shape[1])
        return mean_loss

    def store_transition(self, transition):
        """
        存储经验到回放缓冲区

        参数：
        - transition: 经验元组 (state, action, next_state, reward, done)
        """
        self.replay_buffer.push(*transition)
